package com.hzsparrow.ai.config;

import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.OpenAiEmbeddingOptions;
import org.springframework.stereotype.Component;

import com.hzsparrow.ai.entity.ModelConfig;

import lombok.extern.slf4j.Slf4j;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 嵌入模型管理器
 * 用于动态创建和管理不同类型的嵌入模型
 */
@Slf4j
@Component
public class EmbeddingModelManager {
    
    private final Map<Long, EmbeddingModel> embeddingModelCache = new ConcurrentHashMap<>();
    
    /**
     * 根据模型配置创建嵌入模型
     * 
     * @param modelConfig 模型配置
     * @return 嵌入模型实例
     */
    public EmbeddingModel createEmbeddingModel(ModelConfig modelConfig) {
        if (modelConfig == null) {
            throw new IllegalArgumentException("模型配置不能为空");
        }
        
        if (!"embedding".equals(modelConfig.getModelType())) {
            throw new IllegalArgumentException("模型类型必须是embedding");
        }
        
        try {
            EmbeddingModel embeddingModel = switch (modelConfig.getCategory().toLowerCase()) {
                case "ollama" -> createOllamaEmbeddingModel(modelConfig);
                case "openai" -> createOpenAiEmbeddingModel(modelConfig);
                default -> throw new IllegalArgumentException("不支持的嵌入模型分类: " + modelConfig.getCategory());
            };
            
            log.info("创建嵌入模型成功: {} - {}", modelConfig.getAlias(), modelConfig.getCategory());
            return embeddingModel;
        } catch (Exception e) {
            log.error("创建嵌入模型失败: {} - {}", modelConfig.getAlias(), e.getMessage());
            throw new RuntimeException("创建嵌入模型失败", e);
        }
    }
    
    /**
     * 获取或创建嵌入模型（带缓存）
     * 
     * @param modelConfig 模型配置
     * @return 嵌入模型实例
     */
    public EmbeddingModel getOrCreateEmbeddingModel(ModelConfig modelConfig) {
        return embeddingModelCache.computeIfAbsent(modelConfig.getId(), 
            k -> createEmbeddingModel(modelConfig));
    }
    
    /**
     * 清除指定模型的缓存
     * 
     * @param modelId 模型ID
     */
    public void clearCache(Long modelId) {
        embeddingModelCache.remove(modelId);
        log.info("清除嵌入模型缓存: {}", modelId);
    }
    
    /**
     * 清除所有缓存
     */
    public void clearAllCache() {
        embeddingModelCache.clear();
        log.info("清除所有嵌入模型缓存");
    }
    
    /**
     * 创建Ollama嵌入模型
     */
    private EmbeddingModel createOllamaEmbeddingModel(ModelConfig modelConfig) {
        OllamaApi ollamaApi = OllamaApi.builder()
            .baseUrl(modelConfig.getBasePath())
            .build();
        
        return OllamaEmbeddingModel.builder()
            .ollamaApi(ollamaApi)
            .defaultOptions(OllamaOptions.builder()
                .model(modelConfig.getModelName())
                .build())
            .build();
    }
    
    /**
     * 创建OpenAI嵌入模型
     */
    private EmbeddingModel createOpenAiEmbeddingModel(ModelConfig modelConfig) {
        OpenAiApi openAiApi = OpenAiApi.builder()
            .baseUrl(modelConfig.getBasePath())
            .apiKey(modelConfig.getApiKey())
            .build();
        
        return new OpenAiEmbeddingModel(openAiApi, 
            org.springframework.ai.document.MetadataMode.EMBED,
            OpenAiEmbeddingOptions.builder()
                .model(modelConfig.getModelName())
                .build());
    }
} 